import numpy as np
from shapely.geometry import Polygon, box
import cv2  # 注意：确保你安装了opencv-python

class LesionCrossValidator:
    '''
    用于病灶区域的交叉验证。
    '''
    def __init__(self ,base_image_shape):
        '''
        初始化方法。

        输入:
            lesion_coords_list (list of list of tuples): 每个元素对应一张图像的病灶边界框列表，格式为 (xmin, xmax, ymin, ymax)。
            transform_matrices (list of arrays): 将每张图像映射到基底图像坐标系的转换矩阵。
            base_image_shape (tuple): 基底图像的尺寸（高度, 宽度）。
        '''
        self.base_image_shape = base_image_shape
        # 在基底坐标系中表示的每张图像的覆盖区域和病灶区域
        
    def map_images_to_base(self,transform_matrices):
        '''
        将每张图像的覆盖区域映射到基底坐标系。
        '''
        mapped_image_polygons = [] 
        for M in transform_matrices:
            # 定义图像在自身坐标系中的四个角点
            h, w = self.base_image_shape  # 假设所有图像尺寸一致
            image_corners = np.array([[0, 0],
                                      [w, 0],
                                      [w, h],
                                      [0, h]], dtype=np.float32).reshape(-1, 1, 2)
            # 将角点映射到基底坐标系
            mapped_corners = cv2.perspectiveTransform(image_corners, M)
            # 使用角点生成图像覆盖区域的多边形
            polygon = Polygon(mapped_corners.reshape(-1, 2))
            mapped_image_polygons.append(polygon)
        return mapped_image_polygons
    
    
    def map_lesions_to_base(self,lesion_coords_list, transform_matrices):
        '''
        将每个病灶边界框映射到基底坐标系。
        '''
        mapped_lesion_polygons=[]
        for lesion_bboxes, M in zip(lesion_coords_list, transform_matrices):
            mapped_lesions = []
            for bbox in lesion_bboxes:
                xmin, xmax, ymin, ymax = bbox
                # 定义边界框的四个角点
                lesion_corners = np.array([[xmin, ymin],
                                           [xmax, ymin],
                                           [xmax, ymax],
                                           [xmin, ymax]], dtype=np.float32).reshape(-1, 1, 2)
                # 将角点映射到基底坐标系
                mapped_corners = cv2.perspectiveTransform(lesion_corners, M)
                # 使用角点生成病灶区域的多边形
                polygon = Polygon(mapped_corners.reshape(-1, 2))
                mapped_lesions.append(polygon)
            mapped_lesion_polygons.append(mapped_lesions)
        return mapped_lesion_polygons
    
    def check_image_coverage(self, lesion_polygon, mapped_image_polygons,exclude_idx=None):
        '''
        检查哪些图像覆盖了给定的病灶区域。

        输入:
            lesion_polygon (Polygon): 基底坐标系中的病灶区域。
            exclude_idx (int): 要排除的图像索引（例如，病灶所在图像的索引）。

        返回:
            covered_indices (list of int): 覆盖该病灶区域的图像索引列表。
        '''
        covered_indices = []
        for idx, image_polygon in enumerate(mapped_image_polygons):
            if idx == exclude_idx:
                continue
            if lesion_polygon.intersects(image_polygon):
                covered_indices.append(idx)
        return covered_indices
    
    def compute_overlap(self, polygon1, polygon2):
        '''
        计算两个多边形的重叠面积。

        输入:
            polygon1 (Polygon): 第一个多边形。
            polygon2 (Polygon): 第二个多边形。

        返回:
            overlap_area (float): 重叠区域的面积。
        '''
        intersection = polygon1.intersection(polygon2)
        return intersection.area
    
    def validate_lesions(self,mapped_lesion_polygons,mapped_image_polygons):
        '''
        基于其他图像的病灶重叠情况对病灶进行验证。

        更新:
           validated_lesions: 包含病灶信息的字典列表。
        '''
        validated_lesions=[]
        for idx, lesions in enumerate(mapped_lesion_polygons):
            for lesion_polygon in lesions:
                # 检查其他图像是否覆盖该病灶区域（排除来源图像）
                covered_indices = self.check_image_coverage(lesion_polygon,mapped_image_polygons, exclude_idx=idx)
                
                # 若无其他图像覆盖该区域，病灶为真阳性
                if not covered_indices:
                    validated_lesions.append({
                        'polygon': lesion_polygon,
                        'status': 'True Positive',
                        'images': [idx]
                    })
                    continue  # 跳至下一个病灶
                
                # 检查覆盖图像中的病灶是否重叠
                overlapping_images = []
                for other_idx in covered_indices:
                    other_lesions = mapped_lesion_polygons[other_idx]
                    for other_lesion_polygon in other_lesions:
                        overlap_area = self.compute_overlap(lesion_polygon, other_lesion_polygon)
                        if overlap_area > 0:
                            overlapping_images.append(other_idx)
                            break  # 找到重叠病灶，跳至下一张图像
                
                # 基于重叠病灶判断病灶状态
                if overlapping_images:
                    # 若至少一张图像中存在重叠病灶
                    images_with_lesion = [idx] + overlapping_images
                    validated_lesions.append({
                        'polygon': lesion_polygon,
                        'status': 'True Positive',
                        'images': images_with_lesion
                    })
                else:
                    # 可能为假阳性或伪影
                    validated_lesions.append({
                        'polygon': lesion_polygon,
                        'status': 'False Positive',
                        'images': [idx]
                    })
        return validated_lesions
    
    def get_validated_lesions(self,lesion_coords_list ,transform_matrices):
        '''
        返回:
            validated_lesions (list of dict): 包含病灶信息的列表。
        '''
        mapped_image_polygons = self.map_images_to_base(transform_matrices)
        mapped_lesion_polygons = self.map_lesions_to_base(lesion_coords_list, transform_matrices)
        validate_lesions= self.validate_lesions(mapped_lesion_polygons,mapped_image_polygons)
        return validate_lesions,mapped_image_polygons,mapped_lesion_polygons
